package com.cyberscanner;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import okhttp3.*;
import org.yaml.snakeyaml.Yaml;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.function.Consumer;

public class ScanTask implements Runnable {
    private volatile boolean isRunning = true;
    private final File rootDirectory;
    private final boolean isAuditMode;
    private final boolean includeJsFiles;
    private final ProgressCallback progressCallback;
    private final ColoredMessageCallback messageCallback;
    private final Consumer<String> resultCallback;
    private final OkHttpClient httpClient;
    private final ObjectMapper objectMapper;
    private String ollamaHost;
    private String ollamaModel;

    public ScanTask(File rootDirectory, boolean isAuditMode, boolean includeJsFiles,
                   ProgressCallback progressCallback, ColoredMessageCallback messageCallback, Consumer<String> resultCallback) {
        this.rootDirectory = rootDirectory;
        this.isAuditMode = isAuditMode;
        this.includeJsFiles = includeJsFiles;
        this.progressCallback = progressCallback;
        this.messageCallback = messageCallback;
        this.resultCallback = resultCallback;
        this.httpClient = new OkHttpClient.Builder()
            .connectTimeout(30, java.util.concurrent.TimeUnit.SECONDS)
            .writeTimeout(30, java.util.concurrent.TimeUnit.SECONDS)
            .readTimeout(60, java.util.concurrent.TimeUnit.SECONDS)
            .retryOnConnectionFailure(true)
            .build();
        this.objectMapper = new ObjectMapper();
        loadConfig();
    }

    private void loadConfig() {
        try {
            // 首先尝试从jar包同级目录读取config.yaml
            File externalConfig = new File(new File(getClass().getProtectionDomain()
                    .getCodeSource().getLocation().toURI()).getParent(), "config.yaml");
            
            InputStream inputStream;
            if (externalConfig.exists()) {
                inputStream = new FileInputStream(externalConfig);
            } else {
                // 如果外部配置不存在，尝试从jar包内部读取
                inputStream = getClass().getResourceAsStream("/config.yaml");
                if (inputStream == null) {
                    // 如果两个位置都没有配置文件，显示错误对话框
                    javafx.application.Platform.runLater(() -> {
                        javafx.scene.control.Alert alert = new javafx.scene.control.Alert(javafx.scene.control.Alert.AlertType.ERROR);
                        alert.setTitle("配置错误");
                        alert.setHeaderText("找不到配置文件");
                        alert.setContentText(String.format("请确保在程序目录 %s 下存在 config.yaml 配置文件。\n\n" +
                                "配置文件示例内容：\n" +
                                "api:\n" +
                                "  ollama:\n" +
                                "    url: http://localhost:11434/api\n" +
                                "    model: deepseek-coder",
                                externalConfig.getParent()));
                        alert.showAndWait();
                        System.exit(1);
                    });
                    return;
                }
            }
            
            Yaml yaml = new Yaml();
            Map<String, Object> config = yaml.load(inputStream);
            
            Map<String, Object> api = (Map<String, Object>) config.get("api");
            Map<String, Object> ollama = (Map<String, Object>) api.get("ollama");
            String ollamaUrl = (String) ollama.get("url");
            this.ollamaHost = ollamaUrl.split("/api")[0];
            this.ollamaModel = (String) ollama.get("model");
            
            inputStream.close();
        } catch (Exception e) {
            e.printStackTrace();
            // 显示错误对话框
            javafx.application.Platform.runLater(() -> {
                javafx.scene.control.Alert alert = new javafx.scene.control.Alert(javafx.scene.control.Alert.AlertType.ERROR);
                alert.setTitle("配置错误");
                alert.setHeaderText("配置文件读取失败");
                alert.setContentText("请检查config.yaml文件格式是否正确。\n\n错误信息：" + e.getMessage());
                alert.showAndWait();
                System.exit(1);
            });
        }
    }

    public void stop() {
        isRunning = false;
    }

    @Override
    public void run() {
        Map<String, String> filesContent = scanCodeFiles(rootDirectory);
        List<String> results = new ArrayList<>();
        int totalFiles = filesContent.size();  // 使用实际扫描的文件数量
        int processedFiles = 0;
        int detectedFiles = 0;  // 新增：检测到问题的文件数量

        // 初始化进度条
        javafx.application.Platform.runLater(() -> {
            progressCallback.updateProgress(0.0);
        });

        for (Map.Entry<String, String> entry : filesContent.entrySet()) {
            if (!isRunning) {
                messageCallback.updateMessage("⚠️ 扫描任务已中断", javafx.scene.paint.Color.web("#FFB86C"));
                return;
            }
            String filepath = entry.getKey();
            String content = entry.getValue();

            try {
                String filename = new File(filepath).getName();
                messageCallback.updateMessage(
                    isAuditMode ?
                        String.format("🔍 分析 %s... (%d/%d)", filename, processedFiles + 1, totalFiles) :
                        String.format("🕵️ 扫描 %s... (%d/%d)", filename, processedFiles + 1, totalFiles),
                    javafx.scene.paint.Color.DODGERBLUE);

                String prompt = createPrompt(content);
                String response = callOllamaAPI(prompt);
                String result = processResponse(response);
                
                // 如果检测到问题（包含[高危]标记），增加检测文件计数
                if (!isAuditMode && result.contains("[高危]") && !results.stream().anyMatch(r -> r.contains(filepath))) {
                    detectedFiles++;
                }

                // 根据扫描结果中的风险等级设置不同的颜色
                String formattedResult = String.format("%s %s\n%s\n%s",
                        isAuditMode ? "📄" : "📁",
                        filepath,
                        result,
                        String.join("", Collections.nCopies(50, "━")));
                
                // 根据结果内容设置不同的颜色
                javafx.scene.paint.Color messageColor;
                if (result.contains("[高危]")) {
                    messageColor = javafx.scene.paint.Color.web("#FF4444");
                } else if (result.contains("[中危]")) {
                    messageColor = javafx.scene.paint.Color.web("#FFB86C");
                } else if (result.contains("[低危]")) {
                    messageColor = javafx.scene.paint.Color.web("#50FA7B");
                } else {
                    messageColor = javafx.scene.paint.Color.web("#8BE9FD");
                }
                
                messageCallback.updateMessage(formattedResult, messageColor);
                results.add(formattedResult);

                // 更新进度
                processedFiles++;
                double progress = (double) processedFiles / totalFiles;
                javafx.application.Platform.runLater(() -> {
                    progressCallback.updateProgress(progress);
                });
            } catch (Exception e) {
                String errorMessage = String.format("❌ 错误：%s\n%s", filepath, e.getMessage());
                messageCallback.updateMessage(errorMessage, javafx.scene.paint.Color.web("#6272A4"));
                results.add(errorMessage);
            }
        }

        // Webshell检测模式下不在这里添加统计信息，统计信息由UI层处理
        resultCallback.accept(String.join("\n", results));
    }

    private Map<String, String> scanCodeFiles(File directory) {
        Map<String, String> codeFiles = new HashMap<>();
        List<String> allowedExt = new ArrayList<>(Arrays.asList(
                ".php", ".jsp", ".jspx", ".asp", ".aspx", ".js", ".html", ".py", ".java"
        ));

        if (!includeJsFiles) {
            allowedExt.remove(".js");
            allowedExt.remove(".html");
        }

        scanDirectory(directory, codeFiles, allowedExt);
        return codeFiles;
    }

    private void scanDirectory(File directory, Map<String, String> codeFiles, List<String> allowedExt) {
        File[] files = directory.listFiles();
        if (files != null) {
            for (File file : files) {
                if (file.isDirectory()) {
                    scanDirectory(file, codeFiles, allowedExt);
                } else {
                    String extension = getFileExtension(file);
                    if (allowedExt.contains(extension.toLowerCase())) {
                        try {
                            String content = readFile(file);
                            // 使用相对路径存储文件
                            String relativePath = rootDirectory.toPath().relativize(file.toPath()).toString();
                            codeFiles.put(relativePath, content);
                        } catch (IOException e) {
                            String relativePath = rootDirectory.toPath().relativize(file.toPath()).toString();
                            codeFiles.put(relativePath, "无法读取文件内容");
                        }
                    }
                }
            }
        }
    }

    private String getFileExtension(File file) {
        String name = file.getName();
        int lastIndexOf = name.lastIndexOf(".");
        return lastIndexOf == -1 ? "" : name.substring(lastIndexOf);
    }

    private String readFile(File file) throws IOException {
        StringBuilder content = new StringBuilder();
        try (BufferedReader reader = new BufferedReader(
                new InputStreamReader(new FileInputStream(file), StandardCharsets.UTF_8))) {
            String line;
            while ((line = reader.readLine()) != null) {
                content.append(line).append("\n");
            }
        }
        return content.toString();
    }

    private String createPrompt(String content) {
        if (isAuditMode) {
            return String.format("【强制指令】你是一个专业的安全审计AI，请按以下要求分析代码：\n\n" +
                    "1. 漏洞分析流程：\n" +
                    "   1.1 识别潜在风险点（SQL操作、文件操作、用户输入点、文件上传漏洞、CSRF、SSRF、XSS、RCE、OWASP top10等漏洞）\n" +
                    "   1.2 验证漏洞可利用性\n" +
                    "   1.3 按CVSS评分标准评估风险等级\n\n" +
                    "2. 输出规则：\n" +
                    "   - 仅输出确认存在的高危/中危漏洞\n" +
                    "   - 使用严格格式：[风险等级] 类型 - 位置:行号 - 50字内描述\n" +
                    "   - 禁止解释漏洞原理\n" +
                    "   - 禁止给出修复建议\n" +
                    "   - 如果有可能，给出POC（HTTP请求数据包）\n\n" +
                    "3. 输出示例（除此外不要有任何输出）：\n" +
                    "   [高危] SQL注入 - user_login.php:32 - 未过滤的$_GET参数直接拼接SQL查询\n" +
                    "   [POC]\nPOST /login.php HTTP/1.1\n" +
                    "   Host: example.com\n" +
                    "   Content-Type: application/x-www-form-urlencoded\n" +
                    "   [中危] XSS - comment.jsp:15 - 未转义的userInput输出到HTML\n" +
                    "   [POC]\nPOST /login.php HTTP/1.1\n" +
                    "   Host: example.com\n" +
                    "   Content-Type: application/x-www-form-urlencoded\n\n" +
                    "4. 当前代码（仅限分析）：\n%s", content.substring(0, Math.min(content.length(), 3000)));
        } else {
            return String.format("【Webshell检测指令】请严格按以下步骤分析代码：\n\n" +
                    "1. 检测要求：         \n" +
                    "    请分析以下文件内容是否为WebShell或内存马。要求：\n" +
                    "    1. 检查PHP/JSP/ASP等WebShell特征（如加密函数、执行系统命令、文件操作）\n" +
                    "    2. 识别内存马特征（如无文件落地、进程注入、异常网络连接）\n" +
                    "    3. 分析代码中的可疑功能（如命令执行、文件上传、信息收集）\n" +
                    "    4. 检查混淆编码、加密手段等规避技术\n\n" +
                    "2. 判断规则：\n" +
                    "   - 仅当确认恶意性时报告\n" +
                    "   - 输出格式：🔴 [高危] Webshell - 文件名:行号 - 检测到[特征1+特征2+...]\n\n" +
                    "3. 输出示例（严格按照此格式输出，不要有任何的补充，如果未检测到危险，则不输出，除此之外，不要有任何输出）：\n" +
                    "   🔴 [高危] Webshell - malicious.php:8 - 检测到[system执行+base64解码+错误抑制]\n\n" +
                    "4. 待分析代码：\n%s", content.substring(0, Math.min(content.length(), 3000)));
        }
    }

    private String callOllamaAPI(String prompt) throws IOException {
        MediaType JSON = MediaType.get("application/json; charset=utf-8");
        Map<String, Object> requestMap = new HashMap<>();
        requestMap.put("model", ollamaModel);
        requestMap.put("prompt", prompt);
        requestMap.put("stream", false);
        RequestBody body = RequestBody.create(JSON, objectMapper.writeValueAsString(requestMap));

        Request request = new Request.Builder()
                .url(ollamaHost + "/api/generate")
                .post(body)
                .build();

        try (Response response = httpClient.newCall(request).execute()) {
            if (!response.isSuccessful()) throw new IOException("请求失败: " + response);
            JsonNode jsonResponse = objectMapper.readTree(response.body().string());
            return jsonResponse.get("response").asText();
        }
    }

    private String processResponse(String response) {
        return response.replaceAll("<think>.*?</think>", "");
    }
}